import numpy as np
import pygame
import pygame.gfxdraw
from numba import njit, float64

# 初始化Pygame
pygame.init()

# 屏幕设置
SCREEN_WIDTH, SCREEN_HEIGHT = 1200, 800
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption('带有空气阻力的球在圆形边界内的运动')

# 字体设置
FONT_NAME = 'msyh' if 'msyh' in pygame.font.get_fonts() else 'SimHei'
font = pygame.font.SysFont(FONT_NAME, 16)

# 物理常量
GRAVITY = np.array([0, -9.8])
AIR_RESISTANCE_COEFFICIENT = 0.00  # 空气阻力系数
VELOCITY_POWER = 1  # 空气阻力与速度的次方关系
RESTITUTION_COEFFICIENT = 0.9  # 弹性系数(反弹后的速度等于原来的速度乘以弹性系数)
DT = 0.1  # 时间步长
FPS = 1000  # 帧率
BOUNDARY_RADIUS = 350  # 边界半径

@njit(float64[:](float64[:], float64[:], float64, float64, float64))
def calculate_force(velocity, gravity, air_resistance_coefficient, velocity_power, mass):
    speed = np.linalg.norm(velocity)
    drag_force = -air_resistance_coefficient * speed ** velocity_power * velocity / speed
    total_force = -mass * gravity + drag_force
    return total_force

@njit((float64[:], float64[:], float64[:], float64, float64, float64, float64))
def update_position(position, velocity, gravity, air_resistance_coefficient, velocity_power, dt, mass):
    k1_v = calculate_force(velocity, gravity, air_resistance_coefficient, velocity_power, mass) / mass * dt
    k1_x = velocity * dt

    velocity += 0.5 * k1_v
    k2_v = calculate_force(velocity, gravity, air_resistance_coefficient, velocity_power, mass) / mass * dt
    k2_x = velocity * dt

    velocity = velocity - 0.5 * k1_v + 0.5 * k2_v
    k3_v = calculate_force(velocity, gravity, air_resistance_coefficient, velocity_power, mass) / mass * dt
    k3_x = velocity * dt

    velocity = velocity - 0.5 * k2_v + k3_v
    k4_v = calculate_force(velocity, gravity, air_resistance_coefficient, velocity_power, mass) / mass * dt
    k4_x = velocity * dt

    velocity += (k1_v + 2 * k2_v + 2 * k3_v + k4_v) / 6
    position += (k1_x + 2 * k2_x + 2 * k3_x + k4_x) / 6

    return position, velocity

class Ball:
    """代表一个运动中的球"""
    def __init__(self, position, velocity, radius, color, mass):
        self.position = np.array(position, dtype=float)
        self.velocity = np.array(velocity, dtype=float)
        self.radius = radius
        self.color = color
        self.mass = mass
        self.trail = []  # 轨迹
        self.max_trail_length = 1000  # 最大轨迹长度

    def update_position(self, dt):
        """更新球的位置和速度"""
        self.position, self.velocity = update_position(
            self.position, self.velocity, GRAVITY, AIR_RESISTANCE_COEFFICIENT, VELOCITY_POWER, dt, self.mass
        )

    def check_boundary(self):
        """检查球是否碰到边界，如果碰到则反弹"""
        distance = np.linalg.norm(self.position)
        if distance + self.radius > BOUNDARY_RADIUS:
            direction = self.position / distance
            self.velocity = (self.velocity - 2 * np.dot(self.velocity, direction) * direction) * RESTITUTION_COEFFICIENT
            overlap = distance + self.radius - BOUNDARY_RADIUS
            self.position = direction * (BOUNDARY_RADIUS - self.radius - overlap)

    def update_trail(self):
        """更新球的轨迹"""
        self.trail.append(self.position.copy())
        if len(self.trail) > self.max_trail_length:
            self.trail.pop(0)

    def draw(self, screen, index):
        """在屏幕上绘制球及其轨迹"""
        screen_pos = self.position + [SCREEN_WIDTH // 3, SCREEN_HEIGHT // 2]
        pygame.gfxdraw.aacircle(screen, int(screen_pos[0]), int(screen_pos[1]), self.radius, self.color)
        pygame.gfxdraw.filled_circle(screen, int(screen_pos[0]), int(screen_pos[1]), self.radius, self.color)

        for pos in self.trail:
            screen_trail_pos = pos + [SCREEN_WIDTH // 3, SCREEN_HEIGHT // 2]
            pygame.gfxdraw.aacircle(screen, int(screen_trail_pos[0]), int(screen_trail_pos[1]), 2, self.color)
            pygame.gfxdraw.filled_circle(screen, int(screen_trail_pos[0]), int(screen_trail_pos[1]), 2, self.color)

        self.display_info(screen, index)

    def display_info(self, screen, index):
        """在屏幕上显示球的信息"""
        ke, pe, me = self.get_energies()
        vx, vy = self.velocity
        speed = np.linalg.norm(self.velocity)

        info_text = [
            f"球 {index + 1}:",
            f"速度: {speed:.4f}",
            f"动能: {ke:.4f}",
            f"势能: {pe:.4f}",
            f"机械能: {me:.4f}",
            f"速度: ({vx:.4f}, {vy:.4f})"
        ]

        x_position = 750
        y_offset = 50 + index * 180
        if y_offset + len(info_text) * 20 > SCREEN_HEIGHT:
            y_offset = 50

        for line in info_text:
            text_surface = font.render(line, True, (0, 0, 0))
            screen.blit(text_surface, (x_position, y_offset))
            y_offset += 20

    def get_energies(self):
        """计算球的动能、势能和机械能"""
        ke = 0.5 * self.mass * np.sum(self.velocity ** 2)
        h = self.position[1]
        pe = self.mass * GRAVITY[1] * (h - BOUNDARY_RADIUS)
        return ke, pe, ke + pe

def initialize_balls():
    """初始化球的列表"""
    return [
        Ball(position=[0, 0], velocity=[10 + 1e-9, 0], radius=10, color=(255, 0, 0), mass=1),
        Ball(position=[0, 0], velocity=[10 - 1e-9, 0], radius=10, color=(0, 0, 255), mass=1)
    ]

def main():
    balls = initialize_balls()
    running = True
    clock = pygame.time.Clock()

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        for ball in balls:
            ball.update_position(DT)
            ball.check_boundary()
            ball.update_trail()

        screen.fill((255, 255, 255))
        pygame.gfxdraw.aacircle(screen, SCREEN_WIDTH // 3, SCREEN_HEIGHT // 2, BOUNDARY_RADIUS, (0, 0, 0))
        pygame.gfxdraw.filled_circle(screen, SCREEN_WIDTH // 3, SCREEN_HEIGHT // 2, BOUNDARY_RADIUS, (255, 255, 255))
        pygame.gfxdraw.aacircle(screen, SCREEN_WIDTH // 3, SCREEN_HEIGHT // 2, BOUNDARY_RADIUS, (0, 0, 0))

        for index, ball in enumerate(balls):
            ball.draw(screen, index)

        pygame.display.flip()
        clock.tick(FPS)

    pygame.quit()

if __name__ == "__main__":
    main()